import os
import shutil
import magic_pdf.model as model_config
from loguru import logger
from pathlib import Path
from magic_pdf.pipe.UNIPipe import UNIPipe
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter

"""
解压pdf文件
返回解压后的文件夹路径
"""

def extract(file):
    model_config.__use_inside_model__ = True
    try:
        if not os.path.exists("output"):
            os.mkdir("output")
        demo_name = Path(os.path.basename(file)).stem
        output_dir_path = "output/%s" % demo_name
        if os.path.exists(output_dir_path):
            shutil.rmtree(output_dir_path)
        os.mkdir(output_dir_path)
        output_text_path = "output/%s/text" % demo_name
        pdf_bytes = open(file, "rb").read()
        model_json = []
        jso_useful_key = {"_pdf_type": "", "model_list": model_json}
        image_writer = DiskReaderWriter(output_dir_path)
        pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer)
        pipe.pipe_classify()
        if len(model_json) == 0:
            if model_config.__use_inside_model__:
                pipe.pipe_analyze()
            else:
                logger.error("need model list input")
                exit(1)
        pipe.pipe_parse()
        md_content = pipe.pipe_mk_markdown(output_dir_path, drop_mode="none")
        with open(f"{output_text_path}.md", "w", encoding="utf-8") as f:
            f.write(md_content)
        return output_dir_path

    except Exception as e:
        logger.exception(e)


def extract_all(inout_dictionary):
    files = [file for file in os.listdir(inout_dictionary)]
    for file in files:
        extract(file)
